from langchain.prompts import PromptTemplate
from pydantic import BaseModel, Field
from datetime import datetime
import json
from LLM_Prompt.siliconflow import SiliconFlow

# 使用 SiliconFlow 或 OpenAI 接入大模型
llm = SiliconFlow()  # 替换成你实际的 SiliconFlow 模型接口

classification_prompt = PromptTemplate(
    input_variables=["question", "current_date", "user_questions"],
    template="""  
# 任务说明

你负责根据用户输入的问题，从下方 API 中选择最相关的一个，并准确提取参数。

## 可用 API
1. `vessel_trajectory`:
   - **用途**：获取船只的航线轨迹数据。
   - **必需参数**：
     - `boatName`（船只名称）
     - `startTime`（查询起始时间，ISO8601 格式）
     - `endTime`（查询结束时间，ISO8601 格式）

2. `checkpoint_passage`:
   - **用途**：监控船只经过卡口或检查点的情况。
   - **必需参数**：
     - `boatName`（船只名称）
     - `startTime`（查询起始时间，ISO8601 格式）
     - `endTime`（查询结束时间，ISO8601 格式）

3. `violation_detection`:
   - **用途**：查询船只的违规行为或预警。
   - **必需参数**：
     - `boatName`（船只名称）
     - `startTime`（查询起始时间，ISO8601 格式）
     - `endTime`（查询结束时间，ISO8601 格式）

4. `checkpoint_traffic`:
   - **用途**：查询特定卡口通过的各种船舶流量信息。
   - **必需参数**：
     - `checkpointName`（卡口名称）
         - checkpointName：如涉及卡口，请（模糊）匹配如下：杭申线卡口、长湖申线卡口、金汇港(北闸)、油墩港(油墩港船闸)、苏申内港线(上港宜东)、吴淞江(泗江口站)、黄浦江(松浦大桥南岸)、龙泉港(叶榭塘船闸)、油墩港(东大盈船闸)、大治河船闸、苏申外港线省际检查站、杨思船闸口、东沟船闸口、苏申内港线省际检查站、平申线省际检查站
     - `startTime`（查询起始时间，ISO8601 格式）
     - `endTime`（查询结束时间，ISO8601 格式）

5.与航道有关接口如下：
    5.1 `spatial_computing_channel`:
       - **用途**：查询特定航道的船舶情况。
       - **必需参数**：
         - `waterwayName`（航道名称）
             - waterwayName：如有，提取航道名, 匹配如下：
        "新泾河", "西泗塘", "俞塘", "女儿泾", "大治河", "淀浦河", "六磊塘", "娄塘河", "郭泽塘", "顾浦",
        "孙浜", "蒲汇塘", "周浦塘", "练祁河", "苏申外港线", "苏申内港线", "老杭申线", "老长湖申线",
        "平申线", "杭申线", "长湖申线（太浦河)", "盐铁塘", "蒲华塘", "浏河", "吴淞江", "东高泾",
        "西虬江", "桃浦", "新泾港", "新民港", "创建港", "马家港", "马路河", "走马塘", "湄浦",
        "嘉定城河(东段)", "嘉定城河(西段)", "祁迁河", "横沥", "中槎浦", "漳浦", "浦东运河",
        "随塘河", "咸塘港", "横沔港", "龙泉港", "叶榭塘", "新东海港", "紫石泾", "南泖泾",
        "上横泾", "茹塘", "七仙泾","黄浦江", "油墩港", "华田泾", "新通波塘", "三官塘",
        "航塘港", "奉新港", "虬江", "春申塘", "龙华港", "北沙港", "封浜", "蟠龙塘", "八灶港",
        "木渎港", "黄泥泾", "赵家沟", "川杨河", "高桥港", "外环运河", "新张泾", "池泾",
        "黄姑塘", "定光塘", "老龙泉港", "亭林市河", "新泾塘", "后岗塘", "牛桥港", "白莲泾",
        "张家浜", "姚家浜", "北横泾", "盐仓浦", "潘泾", "北泗塘", "杨盛河", "荻泾", "南沙港",
        "三团港", "南竹港", "大泐港", "古浦塘", "长兴港", "横山塘", "马山塘", "大塔浜",
        "北泖泾", "南湾港", "五厍港", "北石港", "向荡港", "沈泾塘", "通波塘", "大涨泾",
        "人民河", "泗泾塘", "洞泾港", "砖新河", "南界泾", "辰山塘", "昆岗河", "新浜市河",
        "秀春塘", "淀山湖北航道", "毛河泾", "南横港", "北横港", "莲横港", "大港", "东塘港",
        "泖阳港", "汪洋港", "淀山港", "朱昆河", "朱泖河", "洋泾港", "西沟港", "曹家沟",
        "三林塘港", "合庆港", "蒲泽塘", "秀州塘", "白牛塘", "张泾河", "中运河", "东海港",
        "枫泾塘", "惠高泾", "金卫城河", "六里塘", "枫泾市河", "面丈港", "山塘", "二灶港",
        "石皮泐港", "中港河", "六灶港", "三六灶港", "马万港", "盐船港", "新场港", "黄路港",
        "立新港", "小洼港", "金汇港", "浦南运河", "洪庙港", "南横引河", "堡镇港", "六滧港",
        "鸽龙港", "奚家港", "新河港", "东平河", "老滧港", "三沙洪", "新建港", "张网港",
        "四滧港", "八滧港", "庙港", "界河", "西大盈港", "东大盈港", "柘泽塘", "石家浜",
        "老通波塘", "徐泾江", "北中新河", "五尺沟", "惠新港", "泐马河", "清运河", "军民河",
        "惠南城河", "团芦港", "巨潮港", "南横泾", "护塘港", "四团港", "泾罗港", "东沿塘河"
         - `startTime`（查询起始时间，ISO8601 格式）
         - `endTime`（查询结束时间，ISO8601 格式）
     5.2 `spatial_computing_boat`:
       - **用途**：查询船舶在航道上的航行情况。
       - **必需参数**：
         - `boatName`（船只名称）
         - `startTime`（查询起始时间，ISO8601 格式）
         - `endTime`（查询结束时间，ISO8601 格式）

6. `basic_information`:
   - **用途**：查询某条船的基本信息。
   - **必需参数**：
     - `boatName`（船只名称）

## 特殊要求

1. **分类**：对用户输入的问题进行意图分类，只能选择一个最相关 API。
2. **参数提取**：
    - boatName：提取船名（如有）。
    - startTime / endTime：需为 ISO8601 格式，例如 `2025-04-09 00:00:00`。如未指定时间则返回 "未知"。遇“今天”、“昨天”、“前天”等请按当前日期 `{current_date}` 计算。
3. **统计性问题（如总数、统计、汇总等）统一返回 `not matched` API。**
4. **只返回一个最合适的 API 及其参数**
5. **如果{user_questions}不为空，主要根据{question}来提取参数和确定是哪个api，从{user_questions}中的一个或者多个历史问题中提取补充此api需要的其它参数。

## 输出格式

仅返回以下 JSON 结构：

```json
{{  
  "api": "API 名称",  
  "parameters": {{  
    "boatName": "示例船只名称或 null",  
    "startTime": "2025-04-08 00:00:00 或 null",  
    "endTime": "2025-04-08 23:59:59 或 null",  
    "checkpointName": "卡口名称或 null",  
    "waterwayName": "航道名称或 null"  
  }}  
}} 
```

### 示例  
输入  
问题：{question}  
历史问题：{user_questions}

输出  
严格返回上述 JSON 格式内容。

### 说明

- 请勿输出除 JSON 外的任何内容。
- 如匹配多个 API，仅选择并返回你认为最合适的一个。
- 未提及参数填 null。  
- 时间相关按规则转换，无则填 null。
"""
)

# Step 4: Pydantic 数据验证器
class ExtractedData(BaseModel):
    boat_name: str = Field(..., description="船的名字")
    date: str = Field(..., description="日期，格式为 YYYY-MM-DD")

    @staticmethod
    def validate_output(response: str):
        try:
            data = ExtractedData.model_validate_json(response)  # 替换 parse_raw
            return data.model_dump()  # 替换 dict
        except Exception as e:
            raise ValueError(f"数据解析错误: {e}")


# 清理模型输出工具
def clean_model_output(raw_output: str) -> str:
    """
    清理大模型的输出，去除额外的 Markdown 等格式符号，确保输出是合法 JSON。
    """
    # 去掉 Markdown 格式的代码块标记（如 ```json ... ```)
    if raw_output.startswith("```json"):
        raw_output = raw_output.strip("```json").strip("`\n")
    elif raw_output.startswith("```"):
        raw_output = raw_output.strip("```").strip("`\n")

    try:
        # 确保最终输出为合法的 JSON 格式（如有多余的空格，再清理一下）
        json.loads(raw_output)  # 测试是否是有效 JSON
    except json.JSONDecodeError as e:
        raise ValueError(f"清理后的输出仍不是有效 JSON: {e}")

    return raw_output



# Step 5: 链式操作
classification_chain = classification_prompt | llm


def classify_and_extract(question, user_questions):
    # 动态获取当前日期
    current_date = datetime.now().strftime("%Y-%m-%d")

    # 1. 问题分类 提取关键信息
    classification_extract_json = classification_chain.invoke({"question": question, "current_date": current_date, "user_questions": user_questions})

    # 清理模型输出
    cleaned_result = clean_model_output(classification_extract_json)

    return cleaned_result


# 测试代码
# test_queries = [
        # "莲花号前天的航道轨迹",
        # "莲花号2025年1月20号的航道轨迹",
        # "船舶 ABC123 2025年2月20日的行驶轨迹？",
        # "船舶 XYZ789 昨天经过那些卡口？",
        # "船舶 DEF456 本周是否有违规行为（预警）？",
        # "船舶 GHI789 上周去过哪些码头？去过哪些区域？途径那些船闸？途径那些航道？",
        # "卡口 南京长江大桥 2025年2月20日到22日有哪些船舶经过？有哪些违规船舶？有哪些集装箱船舶经过等",
        # "航道 长江主航道 当前有哪些船舶？有哪些违规船舶？有哪些类别船舶？",
        # "2025年2月20日至22日集装箱船的行驶轨迹？",
        # "昨天有哪些集装箱船经过南京长江大桥卡口？",
        # "本周有没有集装箱船存在违规行为？",
        # "长江主航道当前有哪些集装箱船？",
        # "上海2025年1月10日天气怎么样",
        # "查询系统中各类型预警的发生次数和占比，按发生次数从高到低排序。",
        # "江海货368 2025年2月20日航行记录。",
        # "船舶 江海货1819 2024年10月1日到11月30号经过那些卡口？有哪些违规的行为？",
        # "分析2025年1月油墩港(东大盈船闸)每天的流量组成",
        # "查询系统中各类型预警的发生次数和占比，按发生次数从高到低排序。",
        # "统计一月份各卡口点的船舶通过数量，按照通过船舶数量降序排列，只显示通过量前10的卡口。"
# ]

# for example_question in test_queries:
#   print('example_question:', example_question)
#   result = classify_and_extract(example_question)
#   print(result)

# question = "黄浦江航道昨天经过多少条船舶"
# question = "江海货368昨天经过了哪些航道"
# user_questions = []
# result = classify_and_extract(question, user_questions)
# print(result)